#a simple network-based SIR model written in Python
#Jon Zelner
#University of Michigan
#August 11, 2009

#for networks
import igraph

#for random numbers
import random

#what it sounds like; for making copies of objects
import copy

#for plotting
import pylab as pl

#for grabbing command line arguments and doing other 'sys'temsy things
import sys

#check to see if there are command-line arguments; if so bring them in
if len(sys.argv) > 1:
    b = float(sys.argv[1])
    g = float(sys.argv[2])
    N = int(sys.argv[3])
    p = float(sys.argv[4])

#otherwise use defaults
else:
    b = .3
    g = .3
    N = 500
    p = 0.2

neighbors = 2
t = 0

#generate a small-world network
graph = igraph.Graph.Watts_Strogatz(1, N, nei=neighbors, p = p)

#this gets rid of multiple edges and self-loops
graph.simplify()

#now we're going to start tracking connections between individuals. Entries are indexed by individual id. 
#At each index is a list with that individual's connections
adjacencyList = []
for i in range(N):
    adjacencyList.append([])

#just telling connections about each other
for edge in graph.es:
    adjacencyList[edge.source].append(edge.target)
    adjacencyList[edge.target].append(edge.source)

#these hold the IDs of individuals in each of the three states
sAgentList = []
iAgentList = []
rAgentList = []

#these are for holding counts of the number of individuals in each state
sList = []
iList = []
rList = []
newIList = []

#make a list of IDs
allAgents = range(N)

#shuffle the list so we grab individuals in no particular order
random.shuffle(allAgents)

#we start out with everyone susceptible
sAgentList = copy.copy(allAgents)

#take one individual and make them infectious - this is our 'index case'
indexCase = sAgentList.pop()
iAgentList.append(indexCase)

#as long as we have infectious individuals, keep looping
while len(iAgentList) > 0:
    #here, we keep track of who is going to be infected at the *end* of the step and who will recover
    tempIAgentList = []
    recoverList = []
    newI = 0
    
    #just look through the list of infected individuals instead of looking @ everyone
    for iAgent in iAgentList:
        #look in that individual's adjacency list
        for agent in adjacencyList[iAgent]:
            #is that agent susceptible?
            if agent in sAgentList:
                #if so, try to infect
                if (random.random() < b):
                    #if we infect, we can immediately take them off of the susceptible list
                    sAgentList.remove(agent)
                    #but don't add them to the infectious list yet
                    tempIAgentList.append(agent)
                    newI += 1
                    
        #figure out if agent will recover @ end of step
        if (random.random() < g):
            recoverList.append(iAgent)
    
    #now recover agents and remove them from the infectious list
    for recoverAgent in recoverList:
        iAgentList.remove(recoverAgent)
        rAgentList.append(recoverAgent)
    
    #and add newly infected agents to the infectious agent list
    iAgentList.extend(tempIAgentList)
    
    #here, we keep track of how many individuals are in each state
    sList.append(len(sAgentList))
    iList.append(len(iAgentList))
    rList.append(len(sAgentList))
    newIList.append(newI)
    t += 1


for v in graph.vs():
    v['label_size'] = 0
    v['color'] = 'blue'
    if v.index in rAgentList or v.index in iAgentList:
        v['color'] = 'red'
    if v.index == indexCase:
        v['color'] = 'green'

#l = graph.layout_circle()
l = graph.layout_kamada_kawai()
#l = graph.layout_grid_fruchterman_reingold()
igraph.drawing.plot(graph, layout = l)

pl.figure()
pl.plot(iList)
pl.show()

    
            